
#include "histogram_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "register/op_def_registry.h"


namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    HistogramTilingData tiling;
    uint32_t NUM = 8;   //

    uint32_t sizeofdatatype;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();

    auto bins = *context->GetAttrs()->GetInt(0);
    auto min = *context->GetAttrs()->GetFloat(1);
    auto max = *context->GetAttrs()->GetFloat(2);
    if(min == (float)0.0 && max == (float)0.0)
    {
        // printf("SetTilingKey(1)");
        context->SetTilingKey(1);
    }
    else
    {
        // printf("SetTilingKey(2)");
        context->SetTilingKey(2);
    }
        

    auto dt = context->GetInputTensor(0)->GetDataType();
    if(dt == 0)     // float32
    {   if(min == (float)0.0 && max == (float)0.0)
            NUM = 6;
        else
            NUM = 2;
        sizeofdatatype = 4;
    } else if (dt == 1)     //float16
    {   if(min == (float)0.0 && max == (float)0.0)
            NUM = 14;
        else
            NUM = 6;
        sizeofdatatype = 2;
    } else  //int32
    {   if(min == (float)0.0 && max == (float)0.0)
            NUM = 8;
        else
            NUM = 4;
        sizeofdatatype = 4;
    }

    uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;

    uint32_t tiling_size = ((ub_size) / BLOCK_SIZE / 2) / NUM;
    //一次tiling的最大数据
    uint32_t block_size = tiling_size * ALIGN_NUM;

    uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    uint32_t tileNum = totalLengthAligned / block_size;

    uint32_t tileLength = 0;
    uint32_t lasttileLength = 0;
    // 
    if(totalLengthAligned % block_size == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
    {
        if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
        {
            tileNum = 1;
            tileLength = totalLengthAligned;
            lasttileLength = tileLength;
        }
        else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
        {
            tileLength = block_size;
            lasttileLength = tileLength;
        }
        
    }
    else //totalLengthAligned对齐后不能均分到每次tiling
    {
        tileNum = tileNum + 1;
        tileLength = block_size; //tileLength 取最大值(ub)
        lasttileLength = totalLengthAligned % tileLength;
    }


    context->SetBlockDim(1);
    
    tiling.set_ALIGN_NUM(ALIGN_NUM);
    tiling.set_totalLength(totalLength);
    tiling.set_tileNum(tileNum);
    tiling.set_tileLength(tileLength);
    tiling.set_lasttileLength(lasttileLength);

    tiling.set_bins(bins);
    tiling.set_min(min);
    tiling.set_max(max);

    // printf("host bins %d \n", bins);
    // printf("host min %f \n", min);
    // printf("host max %f \n", max);

    
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;   

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Histogram : public OpDef {
public:
    explicit Histogram(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("bins").AttrType(OPTIONAL).Int(100);
        this->Attr("min").AttrType(OPTIONAL).Float(0.0);
        this->Attr("max").AttrType(OPTIONAL).Float(0.0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Histogram);
}
